import dicom
import numpy as np
from plotly.graph_objs import *
from plotly.offline import download_plotlyjs, init_notebook_mode, iplot
init_notebook_mode()

dcm_struct = dicom.read_file("RS1.3.6.1.4.1.2452.6.209502505.1199599685.3559183269.2231927106.dcm", force=True)

structure_names = [item.ROIName for item in dcm_struct.StructureSetROISequence]
structure_names

['Prostate PTV',
 'Urethra',
 'Rectum',
 'Seed Apex',
 'Seed Base',
 'Seed Mid',
 'Avoidance1',
 'Avoidance2',
 'Solid Water',
 'Phantom',
 'Detector']

def pull_structure(number):
    structure_names = [
        item.ROIName for item in dcm_struct.StructureSetROISequence]
    
    contours_by_slice_raw = [
        item.ContourData for item in dcm_struct.ROIContourSequence[number].ContourSequence]
    x = [np.array(item[0::3]) for item in contours_by_slice_raw]
    y = [np.array(item[1::3]) for item in contours_by_slice_raw]
    z = [np.array(item[2::3]) for item in contours_by_slice_raw]
    
    print("Loaded {}".format(structure_names[number]))
    return x, y, z

def display_structures(list_of_structures, colour_list=None):
    dicom_structure_names = np.array(
        [item.ROIName for item in dcm_struct.StructureSetROISequence])
    combined_trace = []
    
    for i, structure in enumerate(list_of_structures):
        if colour_list is None:
            colour = 'black'
        else:
            colour = colour_list[i]
        
        reference = (dicom_structure_names == structure)
        if np.all(reference == False):
            raise Exception("Structure not found (case sensitive)")

        index = int(np.where(reference)[0])
        x, y, z = pull_structure(index)
        
        for i in range(len(x)):
            trace = Scatter3d(
                x=np.append(x[i], x[i][0]), 
                y=np.append(y[i], y[i][0]), 
                z=np.append(z[i], z[i][0]),  
                mode='lines', line=Line(color=colour, width=3))

            combined_trace.append(trace)

    iplot(Figure(
            data=Data(combined_trace),
            layout=Layout(
                showlegend=False,
                width=800, height=1000
            )
        ))

display_structures(
    ['Prostate PTV', 'Urethra', 'Rectum'], 
    ['Red', 'Blue', 'Green']
)

Loaded Prostate PTV
Loaded Urethra
Loaded Rectum

display_structures(
    ['Prostate PTV', 'Urethra', 'Rectum', 'Avoidance1', 'Avoidance2'], 
    ['Red', 'Blue', 'Green', 'Black', 'Black']
)

Loaded Prostate PTV
Loaded Urethra
Loaded Rectum
Loaded Avoidance1
Loaded Avoidance2

display_structures(
    ['Prostate PTV', 'Urethra', 'Rectum', 'Seed Apex', 'Seed Base', 'Seed Mid'], 
    ['Red', 'Blue', 'Green', 'Grey', 'Black', 'Purple']
)

Loaded Prostate PTV
Loaded Urethra
Loaded Rectum
Loaded Seed Apex
Loaded Seed Base
Loaded Seed Mid

display_structures(
    ['Solid Water', 'Phantom', 'Detector'], 
    ['Red', 'Blue', 'Green']
)

Loaded Solid Water
Loaded Phantom
Loaded Detector